from llama_index.embeddings.dashscope import DashScopeEmbedding, DashScopeTextEmbeddingModels
from llama_index.llms.openai import OpenAI
import os
from dotenv import load_dotenv
from typing import Dict
from llama_index.llms.openai.utils import ALL_AVAILABLE_MODELS, CHAT_MODELS



# DEFINED_MODELS: Dict[str, int] = {
#     "qwen-max": 32768,
#     "qwen-plus": 131072,
#     "qwen-turbo": 1000000
# }

DEFINED_MODELS: Dict[str, int] = {
    "Pro/Qwen/Qwen2.5-7B-Instruct":32000,
    "Qwen/Qwen3-30B-A3B":128000,
    "deepseek-ai/DeepSeek-V3":  64000,
    "internlm/internlm2_5-7b-chat":32000,
    "THUDM/GLM-Z1-9B-0414":128000,
    "Qwen/Qwen3-8B": 128000,  # 响应慢
    "deepseek-ai/DeepSeek-R1": 96000,  # 响应慢
}
ALL_AVAILABLE_MODELS.update(DEFINED_MODELS)
CHAT_MODELS.update(DEFINED_MODELS)

load_dotenv()

def dashscope_llm(model_name) -> OpenAI:
    # return OpenAI(
    #     api_key=os.environ.get("DASHSCOPE_API_KEY"),
    #     model=model_name,
    #     api_base="https://dashscope.aliyuncs.com/compatible-mode/v1"
    # )


    return OpenAI(
        api_key=os.environ.get("API_KEY"),
        model=model_name,
        api_base=os.environ.get("API_BASE")
    )



def dashscope_embed_model() -> DashScopeEmbedding:
    embed_model = DashScopeEmbedding(
        # 你也可以使用阿里云提供的其它embedding模型：https://help.aliyun.com/zh/model-studio/getting-started/models#3383780daf8hw
        model_name=DashScopeTextEmbeddingModels.TEXT_EMBEDDING_V3,
        embed_batch_size = 10,
    )
    # 返回创建好的OllamaEmbedding实例
    return embed_model





